﻿#include "DlibMgr.h"

#include "json.h"
#include "../../tool/Utility.h"

#include <dlib/image_io.h>
//#include <facedetectcnn.h>

bool DlibMgr::LoadConfig() {
	Json::Value root;
	if (!Utility::LoadJson("./config/conf_dlib.json", root)) return false;
	this->enable = root["enable"].asBool();
	this->recognition = root["recognition"].asString();
	this->modelconfig  = root["modelconfig"].asString();
	this->resmodelconfig = root["resmodelconfig"].asString();

	return true;
}

int DlibMgr::Init() {
	try
	{
		if (!this->LoadConfig()) {
			return -1;
		}
		if (!this->enable) return 0;
		std::cout << "Starting DlibMgr..." << std::endl;
		// Load face detection and pose estimation models.
		std::cout << dlib::cuda::get_num_devices() << std::endl;
		detector = dlib::get_frontal_face_detector();
		dlib::deserialize(this->modelconfig) >> pose_model;
		dlib::deserialize(this->resmodelconfig) >> anet;
		dlib::deserialize("./models/dlib/mmod_human_face_detector.dat") >> net;

		// 遍历文件夹  1层
		std::vector<cv::String> jpgs;
		cv::glob(this->recognition, jpgs, false);
		std::vector<std::string> facenames;
		std::vector<dlib::matrix<dlib::rgb_pixel>> facedatas;
		for (std::string file_name : jpgs)
		{
			std::cout << file_name << std::endl;
			dlib::matrix<dlib::rgb_pixel> img;
			dlib::load_image(img, file_name);
			for (auto face : detector(img))
			{
				auto shape = pose_model(img, face);
				dlib::matrix<dlib::rgb_pixel> face_chip;
				dlib::extract_image_chip(img, dlib::get_face_chip_details(shape, 150, 0.25), face_chip);
				// TODO 
				facenames.push_back(Utility::GetFileName(file_name));
				facedatas.push_back(std::move(face_chip));
				// 默认每张照片只会出现一个人 其他就不管了
				break;
			}
		}
		auto anet_datas = anet(facedatas);
		for (size_t i =0; i < anet_datas.size(); i++)
		{
			faces_recognition.insert(std::make_pair(facenames[i], anet_datas[i]));
		}
		return 0;
	}
	catch (const std::exception&)
	{
		return -1;
	}
}
unsigned char* pBuffer = (unsigned char*)malloc(0X20000);
std::vector<DlibRet> DlibMgr::Run(cv::Mat& mat) {
	if (!this->enable) return std::vector<DlibRet>(0);

	long time_start = cv::getTickCount();

	dlib::cv_image<dlib::bgr_pixel> cimg(mat);

	std::vector<DlibRet> rets;
	std::vector<dlib::matrix<dlib::rgb_pixel>> refaces;	
	//dlib::matrix<dlib::rgb_pixel> dlibMatrix;
	//dlib::assign_image(dlibMatrix, cimg);

	//for (auto&& d : net(dlibMatrix)) {
	for (auto d : detector(cimg)) {
		auto shape = pose_model(cimg, d);
		DlibRet ret;
		ret.shape = shape;
		rets.push_back(ret);

		dlib::matrix<dlib::rgb_pixel> face_chip;
		dlib::extract_image_chip(cimg, dlib::get_face_chip_details(shape, 150, 0.25), face_chip);
		refaces.push_back(std::move(face_chip));
	}
	
	/* // libfacedetection 版本的人脸检测
	int* pResults = facedetect_cnn(pBuffer, (unsigned char*)(mat.ptr(0)), mat.cols, mat.rows, (int)mat.step);
	std::vector<dlib::rectangle> faces;
	for (int i = 0; i < (pResults ? *pResults : 0); i++)
	{
		short* p = ((short*)(pResults + 1)) + 16 * i;
		int confidence = p[0];
		int x = p[1];
		int y = p[2];
		int w = p[3];
		int h = p[4];
		faces.push_back(dlib::rectangle(x, y, x + w, y + h));
	}*/
	// Find the pose of each face.
	// Detect faces Dlib 官方版本
	//std::vector<dlib::rectangle> faces = detector(cimg);
	//std::vector<DlibRet> rets;
	//std::vector<dlib::matrix<dlib::rgb_pixel>> refaces;
	//for (unsigned long i = 0; i < faces.size(); ++i) {
	//	auto shape = pose_model(cimg, faces[i]);
	//	DlibRet ret;
	//	ret.shape = shape;
	//	rets.push_back(ret);

	//	dlib::matrix<dlib::rgb_pixel> face_chip;
	//	dlib::extract_image_chip(cimg, dlib::get_face_chip_details(shape, 150, 0.25), face_chip);
	//	refaces.push_back(std::move(face_chip));
	//}

	//  TODO 不·能·理·解
	std::vector<dlib::matrix<float, 0, 1>> face_descriptors = anet(refaces);
	
	std::vector<dlib::sample_pair> edges;
	for (size_t i = 0; i < face_descriptors.size(); ++i)
	{
		bool exist = false;
		// 遍历Map
		for (auto& kv: faces_recognition) {
			if (dlib::length(face_descriptors[i] - kv.second) < 0.4) {
				// 嗯 大概就是同一个人吧
				exist = true;
				rets[i].id = kv.first;
				break;
			}
		}
		if (!exist)
		{
			rets[i].id = "Unknown";
		}
	}
	long time_end = cv::getTickCount();
	//std::cout << "DlibMgr Run: " << (time_end - time_start) / cv::getTickFrequency() * 1000 << " " << "ms" << std::endl;
	return rets;
}

cv::Point dlibPoint2cvPoint(const dlib::point &p) {
	return cv::Point(p.x(), p.y());
}

void DlibMgr::Draw(std::vector<DlibRet> rets, const cv::Mat& mat) {
	if (!this->enable) return;
	cv::Scalar color = cv::Scalar(0, 0, 255);
	for (unsigned long i = 0; i < rets.size(); ++i)
	{
		const dlib::full_object_detection& d = rets[i].shape;
		const dlib::rectangle face = d.get_rect();
		// 画 脸部红框
		cv::rectangle(mat, cv::Rect(face.left(), face.top(), face.width(), face.height()), color, 1, 8, 0);
		// 画 文字
		cv::putText(mat, rets[i].id, cv::Point(face.left(), face.top()), cv::FONT_HERSHEY_SIMPLEX, 0.75, cv::Scalar(0, 0, 255), 1);

		// 画 脸部连线

		if (d.num_parts() == 5)
		{
			cv::line(mat, dlibPoint2cvPoint(d.part(0)), dlibPoint2cvPoint(d.part(1)), color);
			cv::line(mat, dlibPoint2cvPoint(d.part(1)), dlibPoint2cvPoint(d.part(4)), color);
			cv::line(mat, dlibPoint2cvPoint(d.part(4)), dlibPoint2cvPoint(d.part(3)), color);
			cv::line(mat, dlibPoint2cvPoint(d.part(3)), dlibPoint2cvPoint(d.part(2)), color);
		}
		else
		{
			// Around Chin. Ear to Ear
			for (unsigned long i = 1; i <= 16; ++i)
				cv::line(mat, dlibPoint2cvPoint(d.part(i)), dlibPoint2cvPoint(d.part(i - 1)), color);

			// Line on top of nose
			for (unsigned long i = 28; i <= 30; ++i)
				cv::line(mat, dlibPoint2cvPoint(d.part(i)), dlibPoint2cvPoint(d.part(i - 1)), color);

			// left eyebrow
			for (unsigned long i = 18; i <= 21; ++i)
				cv::line(mat, dlibPoint2cvPoint(d.part(i)), dlibPoint2cvPoint(d.part(i - 1)), color);
			// Right eyebrow
			for (unsigned long i = 23; i <= 26; ++i)
				cv::line(mat, dlibPoint2cvPoint(d.part(i)), dlibPoint2cvPoint(d.part(i - 1)), color);
			// Bottom part of the nose
			for (unsigned long i = 31; i <= 35; ++i)
				cv::line(mat, dlibPoint2cvPoint(d.part(i)), dlibPoint2cvPoint(d.part(i - 1)), color);
			// Line from the nose to the bottom part above
			cv::line(mat, dlibPoint2cvPoint(d.part(30)), dlibPoint2cvPoint(d.part(35)), color);

			// Left eye
			for (unsigned long i = 37; i <= 41; ++i)
				cv::line(mat, dlibPoint2cvPoint(d.part(i)), dlibPoint2cvPoint(d.part(i - 1)), color);
			cv::line(mat, dlibPoint2cvPoint(d.part(36)), dlibPoint2cvPoint(d.part(41)), color);

			// Right eye
			for (unsigned long i = 43; i <= 47; ++i)
				cv::line(mat, dlibPoint2cvPoint(d.part(i)), dlibPoint2cvPoint(d.part(i - 1)), color);
			cv::line(mat, dlibPoint2cvPoint(d.part(42)), dlibPoint2cvPoint(d.part(47)), color);

			// Lips outer part
			for (unsigned long i = 49; i <= 59; ++i)
				cv::line(mat, dlibPoint2cvPoint(d.part(i)), dlibPoint2cvPoint(d.part(i - 1)), color);
			cv::line(mat, dlibPoint2cvPoint(d.part(48)), dlibPoint2cvPoint(d.part(59)), color);

			// Lips inside part
			for (unsigned long i = 61; i <= 67; ++i)
				cv::line(mat, dlibPoint2cvPoint(d.part(i)), dlibPoint2cvPoint(d.part(i - 1)), color);
			cv::line(mat, dlibPoint2cvPoint(d.part(60)), dlibPoint2cvPoint(d.part(67)), color);
		}
	}
}